#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooUniform.h"
#include "RooAddPdf.h"
#include "RooDataSet.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "RooFitResult.h"
#include "TLegend.h"
#include <iostream>
#include <memory>

using namespace RooFit;

void comparison_test_numerical() {
    RooRealVar x("x", "x", 0, 20);
    RooGaussian signal_truth("signal_truth", "signal_truth", x, RooConst(10.0), RooConst(0.5));
    RooRealVar bkg_low_truth("bkg_low_truth", "bkg_low_truth", 4.0);
    RooRealVar bkg_high_truth("bkg_high_truth", "bkg_high_truth", 16.0);
    RooUniform background_truth("background_truth", "background_truth", x, bkg_low_truth, bkg_high_truth);
    RooRealVar n_sig_truth("n_sig_truth", "n_sig_truth", 300);
    RooRealVar n_bkg_truth("n_bkg_truth", "n_bkg_truth", 700);
    RooAddPdf model_truth("model_truth", "model_truth", {signal_truth, background_truth}, {n_sig_truth, n_bkg_truth});
    std::unique_ptr<RooDataSet> data{model_truth.generate(x, 1000)};
    TCanvas* c = new TCanvas("c", "RooUniform Before vs. After", 1200, 600);
    c->Divide(2, 1);

    c->cd(1);
    RooPlot* frame_before = x.frame(Title("The Problem (Old RooUniform)"));
    data->plotOn(frame_before);
    RooUniform background_old("background_old", "background_old", x);
    RooRealVar n_sig_old("n_sig_old", "n_sig_old", 300, 0, 10000);
    RooRealVar n_bkg_old("n_bkg_old", "n_bkg_old", 700, 0, 10000);
    RooAddPdf model_old("model_old", "model_old", {signal_truth, background_old}, {n_sig_old, n_bkg_old});
    std::unique_ptr<RooFitResult> fitResult_old{model_old.fitTo(*data, Save(), PrintLevel(-1))};
    model_old.plotOn(frame_before, LineColor(kRed), Name("fit_old"));
    frame_before->Draw();

    c->cd(2);
    RooPlot* frame_after = x.frame(Title("The Solution (New RooUniform)"));
    data->plotOn(frame_after);
    RooRealVar bkg_low_fit("bkg_low_fit", "bkg_low_fit", 3.0, 0.0, 8.0);
    RooRealVar bkg_high_fit("bkg_high_fit", "bkg_high_fit", 17.0, 12.0, 20.0);
    RooUniform background_new("background_new", "background_new", x, bkg_low_fit, bkg_high_fit);
    RooRealVar n_sig_new("n_sig_new", "n_sig_new", 300, 0, 10000);
    RooRealVar n_bkg_new("n_bkg_new", "n_bkg_new", 700, 0, 10000);
    RooAddPdf model_new("model_new", "model_new", {signal_truth, background_new}, {n_sig_new, n_bkg_new});
    std::unique_ptr<RooFitResult> fitResult_new{model_new.fitTo(*data, Save(), PrintLevel(-1))};
    model_new.plotOn(frame_after, LineColor(kGreen+2), Name("fit_new"));
    model_new.plotOn(frame_after, Components("background_new"), LineStyle(kDashed), LineColor(kMagenta));
    frame_after->Draw();

    std::cout << "\n\n--- Numerical Comparison ---" << std::endl;
    std::cout << "\n--- Fit with OLD RooUniform (The Problem) ---" << std::endl;
    fitResult_old->Print();
    std::cout << "\n--- Fit with NEW RooUniform (The Solution) ---" << std::endl;
    fitResult_new->Print();

    c->SaveAs("comparison_test_numerical.png");
}